import os
import numpy as np
from PIL import Image
from onnxruntime.quantization import CalibrationDataReader, QuantFormat, quantize_static, QuantType, CalibrationMethod
import onnx

# 模型路径
model_fp32 = 'your_model.onnx'
model_quant_static = 'your_model_quant.onnx'

print("🔍 正在读取模型节点信息...")
onnx_model = onnx.load(model_fp32)

conv_nodes = []
for node in onnx_model.graph.node:
    if node.op_type == 'Conv':
        conv_nodes.append(node.name)
        print(f"发现Conv节点: {node.name}")

print(f"📊 总共发现 {len(conv_nodes)} 个Conv节点")
 
# 数据预处理(根据对应任务修改)
def preprocess_image(srcimg):
    h, w = srcimg.shape[:2]
    inp_height = 48
    inp_width = 2048
    
    # 计算宽高比
    ratio = w / float(h)
    resized_w = int(np.ceil(float(inp_height) * ratio))
    
    # 限制最大宽度
    if np.ceil(inp_height * ratio) > inp_width:
        resized_w = inp_width
    
    pil_img = Image.fromarray(srcimg)
    resized_img = pil_img.resize((resized_w, inp_height), Image.BILINEAR)
    
    # 转换为numpy数组
    img_array = np.asarray(resized_img, dtype='float32')
    
    return normalize_image(img_array, inp_width)

def normalize_image(img, inp_width):
    row, col = img.shape[:2]
    channels = img.shape[2] if len(img.shape) == 3 else 1
    
    resize_w = col
    resize_h = row
    
    # 初始化输出数组
    if channels == 3:
        output = np.zeros((channels, resize_h, inp_width), dtype='float32')
        
        for c in range(3):
            for i in range(row):
                for j in range(inp_width):
                    if j < col:
                        pix = img[i, j, c]
                        output[c, i, j] = (pix / 255.0 - 0.5) / 0.5
                    else:
                        # 填充为0
                        output[c, i, j] = 0.0
        
        # 添加batch维度
        return output[None, ...]
    else:
        # 处理单通道图像
        output = np.zeros((1, resize_h, inp_width), dtype='float32')
        for i in range(row):
            for j in range(inp_width):
                if j < col:
                    pix = img[i, j]
                    output[0, i, j] = (pix / 255.0 - 0.5) / 0.5
                else:
                    output[0, i, j] = 0.0
        
        return output[None, ...]

# 用于校准的图像数据
img_dir = 'your_calibration_image_dir'
img_num = 35
datas = [
    preprocess_image(
        np.asarray(Image.open(os.path.join(img_dir, img)).convert('RGB'))
    ) for img in os.listdir(img_dir)[:img_num]
]
 
# 数据批次读取器
def batch_reader(datas, batch_size):
    _datas = []
    length = len(datas)
    for i, data in enumerate(datas):
        if batch_size==1:
            yield {'x': data}
        elif (i+1) % batch_size==0:
            _datas.append(data)
            yield {'x': np.concatenate(_datas, 0)}
            _datas = []
        elif i<length-1:
            _datas.append(data)
        else:
            _datas.append(data)
            yield {'x': np.concatenate(_datas, 0)}
 
# 构建校准数据读取器
class DataReader(CalibrationDataReader):
    def __init__(self, datas, batch_size):
        self.datas = batch_reader(datas, batch_size)
 
    def get_next(self):
        return next(self.datas, None)
 
# 实例化一个校准数据读取器
data_reader = DataReader(datas, 1)
 
# 静态量化
quantize_static(
    model_input=model_fp32, # 输入模型
    model_output=model_quant_static, # 输出模型
    calibration_data_reader=data_reader, # 校准数据读取器
    quant_format=QuantFormat.QDQ, # 量化格式 QDQ / QOperator
    activation_type=QuantType.QUInt8, # 激活类型 Int8 / UInt8
    weight_type=QuantType.QUInt8, # 参数类型 Int8 / UInt8
    calibrate_method=CalibrationMethod.Percentile, # 数据校准方法 MinMax / Entropy / Percentile
    per_channel=False,  # 是否对每个通道进行量化
    reduce_range=False,  # 是否使用减少的范围
    # nodes_to_quantize=conv_nodes,
    # nodes_to_exclude=['Conv.25', 'Conv.26']
)